#include "stdafx.h"
#include "timer.h"

using namespace concurrency;
using namespace concurrency::fast_math;

int isnan(float f) restrict(cpu)
{
    return _isnan(f);
}

template<typename T,int Rank>
void Fill(array<T,Rank>& arr,T initValue) { 
    parallel_for_each(arr.extent,[&arr,initValue](index<Rank> idx)restrict(amp)
    {
        arr[idx] = initValue; 
    }); 
}
// Core mandelbrot function, returns iteration plus distance as float
float mandelbrot_calc(
    const unsigned max_iterations, 
    const float x0,
    const float y0 
    ) restrict(cpu,amp)
{
    //from http://en.wikipedia.org/wiki/Mandelbrot_set#Optimizations
    //detect main cardoid quickly
    float q0 = x0 - 0.25f; 
    float q = q0*q0 + y0*y0;
    if(q*(q+q0) < y0*y0/4)
    {
        return 0;
    }

    //from http://en.wikipedia.org/wiki/Mandelbrot_set#Continuous_.28smooth.29_coloring
    const float escape = 1e10f;
    float y = 0.0f;
    float x = 0.0f;
    float yy = 0.0f;
    float xx = 0.0f;
    unsigned int iteration = 0;
    while (xx + yy <= escape)
    {
        y = 2.0f * x * y + y0;
        x = xx - yy + x0;
        yy = y * y;
        xx = x * x;
        iteration += 1;
        if(iteration == max_iterations)
            return 0;
    }
    float i = log(log(xx + yy) / log(escape)) / log(2.0f);
    return log(iteration - i);
}

template<unsigned int tile_size>
void brot_amp_optimised_core(
        array<float, 2>& device_dists,
        const float left, 
        const float width, 
        const float top, 
        const float height, 
        const int max_iterations,
        const int oversample)
{
    int init_value = 0;
    array<int, 1> next_chunk(1, &init_value);
    int w=device_dists.extent[0], h=device_dists.extent[1];

    const int max_chunks = (w * h) / (tile_size * tile_size);
    const float xscale = width / (float)w;
    const float yscale = height / (float)h;

    parallel_for_each(
        device_dists.extent.tile<tile_size, tile_size>(),
        [=, &device_dists, &next_chunk] (tiled_index<tile_size, tile_size> idx) restrict(amp)
    {
        tile_static int chunk_id;
        tile_static int global_y;
        tile_static int global_x;

        // Here each tile will process a chuck of data and pick next block to process
        // This is like load balancing computation, a tile will pick next available chunk to process
        // "chunk_id" value in a tile will determine which chunk of data is being processed by this tile
        while (1)
        {
            // All threads from previous iteration sync here
            idx.barrier.wait();
            if (idx.local[0] == 0 && idx.local[1] == 0)
            {
                // Sync-ing chuck to be processed between tiles
                // get chunk to process for this tile
                chunk_id = atomic_fetch_add(&next_chunk[0], 1);
                global_x = chunk_id % (w / tile_size) * tile_size;
                global_y = chunk_id / (w / tile_size) * tile_size;
            }
            // Sync within a tile.
            // Now threads have tile specific chunk_id, global_y, and global_x
            idx.barrier.wait();

            if (chunk_id >= max_chunks) break; 

            // calculate Mandelbrot for scaled coordinate of pixel 
            int pixel_x = global_x + idx.local[0];
            int pixel_y = global_y + idx.local[1];
            float sum = 0.0f;
            for(int i=0; i<oversample; ++i)
                for(int j=0; j<oversample; ++j)
            {
                float x0 = left + (pixel_x+i/oversample) * xscale;
                float y0 = top + (pixel_y+j/oversample) * yscale;
                sum += mandelbrot_calc(max_iterations, x0, y0);
            }
            device_dists(pixel_x, pixel_y) = sum / (oversample*oversample);
        }
    });
}

template<unsigned int tile_size>
void brot_amp_optimised(
    float* dists, 
    const int w, 
    const int h, 
    const float left, 
    const float width, 
    const float top, 
    const float height, 
    const int max_iterations,
    const int oversample)
{
    array<float, 2> device_dists(w, h);
    brot_amp_optimised_core<tile_size>(device_dists, left, width, top, height, max_iterations, oversample);
    array_view<float,2> distsView(w, h, dists);
    copy(device_dists, distsView);
    distsView.synchronize();
}

extern "C" __declspec ( dllexport ) void _stdcall 
    brot_amp_optimised(
    float* dists, 
    const int w, 
    const int h, 
    const float left, 
    const float width, 
    const float top, 
    const float height, 
    const int max_iterations,
    const int oversample)
{
    brot_amp_optimised<16>(dists, w, h, left, width, top, height, max_iterations, oversample);
}

extern "C" __declspec ( dllexport ) void _stdcall buddhabrot_amp(
    float* dists, 
    const int w, 
    const int h,
    const int min_iterations,
    const int max_iterations)
{
    const float left = 0.3033589f, top = 0.4435042f, width=3E-4f, height=3E-4f;

    int init_value = 0;
    array<int, 1> next_chunk(1, &init_value);
    array<float, 2> device_dists(w, h);
    brot_amp_optimised_core<16>(device_dists, left, width, top, height, max_iterations, 1);

    array<float, 2> buddha(w, h);
    Fill(buddha, 0.0f);

    const int tile_size = 4;
    const int max_chunks = (w * h) / (tile_size * tile_size);
    const float xscale = width / (float)w;
    const float yscale = height / (float)h;

    parallel_for_each(
        device_dists.extent.tile<tile_size, tile_size>(),
        [=, &device_dists, &buddha, &next_chunk] (tiled_index<tile_size, tile_size> idx) restrict(amp)
    {
        //tile_static int chunk_id;
        //tile_static int global_y;
        //tile_static int global_x;

        // Here each tile will process a chuck of data and pick next block to process
        // This is like load balancing computation, a tile will pick next available chunk to process
        // "chunk_id" value in a tile will determine which chunk of data is being processed by this tile
        //while (1)
        //{
        //    // All threads from previous iteration sync here
        //    idx.barrier.wait();
        //    if (idx.local[0] == 0 && idx.local[1] == 0)
        //    {
        //        // Sync-ing chuck to be processed between tiles
        //        // get chunk to process for this tile
        //        chunk_id = atomic_fetch_add(&next_chunk[0], 1);
        //        global_x = chunk_id % (w / tile_size) * tile_size;
        //        global_y = chunk_id / (w / tile_size) * tile_size;
        //    }
        //    // Sync within a tile.
        //    // Now threads have tile specific chunk_id, global_y, and global_x
        //    idx.barrier.wait();

        //    if (chunk_id >= max_chunks) break; 
        int global_x = idx.global[0]*tile_size, global_y=idx.global[1]*tile_size;

            int pixel_x = global_x + idx.local[0];
            int pixel_y = global_y + idx.local[1];

            if(device_dists(pixel_x, pixel_y)>log((float)min_iterations)) //only if the point isn't in the Mandelbrot set
            {
                float x0 = left + pixel_x * xscale;
                float y0 = top + pixel_y * yscale;
                float y = 0.0f;
                float x = 0.0f;
                float yy = 0.0f;
                float xx = 0.0f;
                int iteration = 0;
                while (xx + yy < 4 && iteration < max_iterations)
                {
                    y = 2.0f * x * y + y0;
                    x = xx - yy + x0;
                    yy = y * y;
                    xx = x * x;
                    iteration += 1;
                    if(-1<x && x<1 && -1<y && y<1)
                        buddha(int((x+1.0f)*w/2), int((y+1.0f)*h/2))+=1.0f; //increment point
                }
            }
        //}
    });

    array_view<float,2> distsView(w, h, dists);
    copy(buddha, distsView);
    distsView.synchronize();
}


extern "C" __declspec ( dllexport ) void _stdcall 
    brot_amp_basic(float* dists, int w, int h, float left, float width, float top, float height, int max_iterations, const int oversample)
{
    //from http://en.wikipedia.org/wiki/Mandelbrot_set#Continuous_.28smooth.29_coloring

    array<float, 2> device_dists(w, h);

    float yscale = height / (float)h;
    float xscale = width / (float)w;

    parallel_for_each(device_dists.extent, [=, &device_dists] (index<2> idx) restrict(amp)
    {
        // calculate Mandelbrot for scaled coordinate of pixel 
        float x0 = idx[0] * xscale + left;
        float y0 = idx[1] * yscale + top;
        device_dists(idx) = mandelbrot_calc(max_iterations, x0, y0);
    });

    array_view<float,2> distsView(w, h, dists);
    copy(device_dists, distsView);
    distsView.synchronize();
}

extern "C" __declspec ( dllexport ) void _stdcall 
    brot_cpu_single(float* dists, int w, int h, float left, float width, float top, float height, int max_iterations, const int oversample)
{
    for(int pixel_x = 0; pixel_x<w; ++pixel_x)
        for(int pixel_y = 0; pixel_y<h; ++pixel_y)
    {
        // calculate Mandelbrot for scaled coordinate of pixel 
        float x0 = pixel_x * height / h + left;
        float y0 = pixel_y * width / w + top;
        dists[pixel_x*w+pixel_y] = mandelbrot_calc(max_iterations, x0, y0);
    }
}
extern "C" __declspec ( dllexport ) void _stdcall 
    brot_cpu_parallel(float* dists, int w, int h, float left, float width, float top, float height, int max_iterations, const int oversample)
{
    parallel_for(0, h, [=] (int pixel_y) restrict(cpu)
    {
        for(int pixel_x = 0; pixel_x<w; ++pixel_x)
        {
            // calculate Mandelbrot for scaled coordinate of pixel 
            float x0 = pixel_x * height / h + left;
            float y0 = pixel_y * width / w + top;
            dists[pixel_x*w+pixel_y] = mandelbrot_calc(max_iterations, x0, y0);
        }
    });
}

extern "C" __declspec ( dllexport ) void _stdcall 
    get_device_name(wchar_t* buffer)
{
    std::wstring name = accelerator().get_description();
    //copy(name.begin(), name.end(), buffer);
    //buffer[name.size()*sizeof(wchar_t)]=0;
    wcscpy(buffer, name.c_str());
}

void perftest()
{
    unsigned int max_iterations=256;
    for(int size=1<<8; size<=1<<13; size*=2)
    {
        std::vector<float> out(size*size, 0);
        for(int test=0; test<=3; ++test)
        {
            switch (test)
            {
            case 0:
                std::cout << "Basic" << std::endl; break;
            case 1:
                std::cout << "Fast" << std::endl; break;
            case 2:
                std::cout << "CPU singlethread" << std::endl; break;
            case 3:
                std::cout << "CPU parallel" << std::endl; break;
            }
            for(int tile_size=2; tile_size<=(test==1?32:8); tile_size*=2)
                for(int oversample=1; oversample<=4; ++oversample)
            {
                try 
                {
                    Timer time;
                    time.Start();
                    switch (test)
                    {
                        case 0:
                            brot_amp_basic(&out[0], size, size, -2.0f, 4.0f, -2.0f, 4.0f, max_iterations, oversample);
                            break;
                        case 1:
                            switch (tile_size)
                            {
                            case 2: brot_amp_optimised<2>(&out[0], size, size, -2.0f, 4.0f, -2.0f, 4.0f, max_iterations, oversample); break;
                            case 4: brot_amp_optimised<4>(&out[0], size, size, -2.0f, 4.0f, -2.0f, 4.0f, max_iterations, oversample); break;
                            case 8: brot_amp_optimised<8>(&out[0], size, size, -2.0f, 4.0f, -2.0f, 4.0f, max_iterations, oversample); break;
                            case 16: brot_amp_optimised<16>(&out[0], size, size, -2.0f, 4.0f, -2.0f, 4.0f, max_iterations, oversample); break;
                            case 32: brot_amp_optimised<32>(&out[0], size, size, -2.0f, 4.0f, -2.0f, 4.0f, max_iterations, oversample); break;
                            }
                            break;
                        case 2:
                            brot_cpu_single(&out[0], size, size, -2.0f, 4.0f, -2.0f, 4.0f, max_iterations, oversample);
                            break;
                        case 3:
                            brot_cpu_parallel(&out[0], size, size, -2.0f, 4.0f, -2.0f, 4.0f, max_iterations, oversample);
                            break;
                    }
                    time.Stop();
                    std::cout 
                        << "size " << size 
                        << "\ttile " << std::setw(3) << (test==1?tile_size:1)
                        << "\tover " << std::setw(1) << oversample
                        << "\ttime " << std::setw(7) << time.Elapsed() 
                        << "\tMpixels/s " << (size*size*(test==1?oversample*oversample:1))/(time.Elapsed()/1000)/1e6
                        << std::endl;
                }
                catch(std::exception e)
                {
                    std::cout << "EXCEPTION " << e.what() << std::endl;
                }
            }
        }
        std::cout << std::endl;
    }
}

void main()
{
    perftest();
}
